package com.pw.study.flink.app

import com.alibaba.fastjson.JSONObject
import com.pw.study.flink.source.{CustomSourceByRich, StudentSource, UserPO}
import org.apache.flink.api.common.state.{BroadcastState, MapStateDescriptor, ReadOnlyBroadcastState}
import org.apache.flink.streaming.api.TimeCharacteristic
import org.apache.flink.streaming.api.datastream.BroadcastStream
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction
import org.apache.flink.streaming.api.functions.timestamps.BoundedOutOfOrdernessTimestampExtractor
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.function.ProcessWindowFunction
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows
import org.apache.flink.streaming.api.windowing.time.Time
import org.apache.flink.streaming.api.windowing.windows.TimeWindow
import org.apache.flink.util.Collector

object CustomApp {

  def addDim(ds: DataStream[UserPO], dsData: DataStream[JSONObject]): DataStream[JSONObject] = {
    val user: MapStateDescriptor[String, UserPO] = new MapStateDescriptor[String, UserPO]("user", classOf[String], classOf[UserPO])
    val dsUser: BroadcastStream[UserPO] = ds.broadcast(user)
    val stream: DataStream[JSONObject] = dsData.connect(dsUser).process(new BroadcastProcessFunction[JSONObject, UserPO, JSONObject] {
      private var num = 0

      override def processElement(in1: JSONObject, ctx: BroadcastProcessFunction[JSONObject, UserPO, JSONObject]#ReadOnlyContext, collector: Collector[JSONObject]): Unit = {
         println(s"===============维度： ${ctx.currentWatermark()}")
        val stateMap: ReadOnlyBroadcastState[String, UserPO] = ctx.getBroadcastState(user)
        num = num + 1;
        if (num > 5) {
          num = 1
        }
        var po: UserPO = stateMap.get(num.toString)
        if (po == null) {
          po = UserPO("1", "蝎子", 33, "future@123")
        }
        val id: String = po.id
        val name: String = po.name
        val age: Int = po.age
        val email: String = po.email
        val obj = new JSONObject()
        obj.put("id", id)
        obj.put("name", name)
        obj.put("age", age)
        obj.put("email", email)

        in1.put("User", obj)
        collector.collect(in1)
      }

      override def processBroadcastElement(in2: UserPO, ctx: BroadcastProcessFunction[JSONObject, UserPO, JSONObject]#Context, collector: Collector[JSONObject]): Unit = {
        println("输出-----------------------------------------")
        val stateMap: BroadcastState[String, UserPO] = ctx.getBroadcastState(user)
        stateMap.put(in2.id, in2)
      }
    })
    stream
  }

  def addWatermark(dsDim: DataStream[JSONObject]): DataStream[JSONObject] = {
    dsDim.assignTimestampsAndWatermarks(new BoundedOutOfOrdernessTimestampExtractor[JSONObject](Time.seconds(3)) {
      override def extractTimestamp(t: JSONObject): Long = {
        println("----mark: " + getCurrentWatermark)
        t.getLong("ts")
      }
    })
  }

  def addWindow(dsWm: DataStream[JSONObject]): DataStream[JSONObject] = {
    dsWm.keyBy(_.getJSONObject("User").getString("name")).window(TumblingEventTimeWindows.of(Time.seconds(30)))
      //.trigger(new MyTrigger)
      .process(new ProcessWindowFunction[JSONObject, JSONObject, String, TimeWindow] {
        override def process(key: String, ctx: Context, elements: Iterable[JSONObject], out: Collector[JSONObject]): Unit = {
          println(s"======key: $key  window: ${ctx.window.getStart} ==> ${ctx.window.getEnd}  mark: ${ctx.currentWatermark}  ")
          val arr: Array[JSONObject] = elements.toArray.sortBy(_.getLong("ts"))
          for (elem <- arr) {
            out.collect(elem)
          }
        }
      })
  }

  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment
    env.setStreamTimeCharacteristic(TimeCharacteristic.EventTime)
    env.setParallelism(5)
    val ds: DataStream[UserPO] = env.addSource(new CustomSourceByRich())
    val dsData: DataStream[JSONObject] = env.addSource(new StudentSource)
    val dsDim: DataStream[JSONObject] = addDim(ds, dsData)
    val dsWm: DataStream[JSONObject] = addWatermark(dsDim)
    val dsW: DataStream[JSONObject] = addWindow(dsWm)
    dsW.print("窗口：")

    env.execute("RichSource")
  }
}
